# -*- coding: utf-8 -*-
# coding=utf-8

import copy
import re
from utils.log import logger
from transformers import AutoTokenizer, AutoModel
import json

class HeaderRecognition:
    def __init__(self, model, headers_path) -> None:
        self.model = model
        #读取所有候选项
        with open(headers_path, 'r') as file:
            data = json.load(file)
        self.candidates = data
        self.prompt = f"你现在是一个语义分析模型，请你帮我从给定的候选列 [{'，'.join(self.candidates)}]选择文本最有可能涉及的列。"
        self.prompt = self.prompt + "文本：{}\n涉及的列："

    def recognize(self, query):
        """
        识别query对应的header
        :param query: 自然语言查询
        :return:一个header列表
        """
        try:
            prediction = self.model.predict(self.prompt.format(query))
        except Exception as err:
            logger.error("header prediction error happened![{}]: {}".format(type(err), str(err)))
            raise
        headers = self.get_pred_data(prediction)
        headers = self.check_headers(headers)
        return headers

    def get_pred_data(self, origin_data):
        pattern = r'\[(.*?)\]'
        pred_data = re.findall(pattern, origin_data)
        if len(pred_data) > 0:
            pred_data = pred_data[0]
        else:
            return []
        if ',' in pred_data:
            pred_data = pred_data.split(',')
        else:
            if '，' in pred_data:
                pred_data = pred_data.split('，')
        return pred_data

    def check_headers(self, header):
        new_headers = []
        for head in header:
            if head in self.candidates:
                new_headers.append(head)
        return new_headers
        

